#!/usr/bin/env python
# -*- coding: utf-8 -*-


import numpy as np

def driver_noniid(dataset, num_users):
    """
    Sample non-I.I.D client data from MNIST dataset
    :param dataset:
    :param num_users:
    :return:
    """
    # 将数据集分成200个shard，每个shard包含108个图像
    num_shards, num_imgs = 200, 108
    # 将200个shard的索引存入一个列表中
    idx_shard = [i for i in range(num_shards)]
    # 创建一个字典dict_users，其中键为用户索引，值为一个空的Numpy数组。该字典将用于存储每个用户的本地数据集
    dict_users = {i: np.array([], dtype='int64') for i in range(num_users)}
    # 生成一个包含所有图像索引的Numpy数组
    idxs = np.arange(num_shards*num_imgs)
    # 将训练集中的标签转换为Numpy数组
    labels = np.array(dataset.targets)
    # 将图像索引和标签合并成一个矩阵，并按照标签排序。最终得到一个按照标签排序的图像索引数组
    # sort labels
    idxs_labels = np.vstack((idxs, labels))
    idxs_labels = idxs_labels[:,idxs_labels[1,:].argsort()]
    idxs = idxs_labels[0,:]
    # 对于每个用户，从200个shard中随机选择20个shard，
    # 将这两个shard的图像索引合并成一个数组，并将该数组添加到dict_users中对应用户的本地数据集中
    # divide and assign
    for i in range(num_users):
        rand_set = set(np.random.choice(idx_shard, 20, replace=False))
        idx_shard = list(set(idx_shard) - rand_set)
        for rand in rand_set:
            dict_users[i] = np.concatenate((dict_users[i], idxs[rand*num_imgs:(rand+1)*num_imgs]), axis=0)
    return dict_users

